{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZloPIuRHn97X"
      },
      "source": [
        "##### Copyright 2018 The TensorFlow Authors. [Licensed under the Apache License, Version 2.0](#scrollTo=Afd8bu4xJOgh)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tNgCmfUvJNoF"
      },
      "outputs": [],
      "source": [
        "// #@title Licensed under the Apache License, Version 2.0 (the \"License\"); { display-mode: \"form\" }\n",
        "// Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "// you may not use this file except in compliance with the License.\n",
        "// You may obtain a copy of the License at\n",
        "//\n",
        "// https://www.apache.org/licenses/LICENSE-2.0\n",
        "//\n",
        "// Unless required by applicable law or agreed to in writing, software\n",
        "// distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "// See the License for the specific language governing permissions and\n",
        "// limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AlvdCHw5JGyx"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td><a target=\"_blank\" href=\"https://www.tensorflow.org/swift/tutorials/custom_differentiation\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\">}TensorFlow.org에서 보기</a></td>\n",
        "  <td><a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/swift/blob/master/docs/site/tutorials/custom_differentiation.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\">Run in Google Colab</a></td>\n",
        "  <td><a target=\"_blank\" href=\"https://github.com/tensorflow/swift/blob/master/docs/site/tutorials/custom_differentiation.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\">GitHub에서소스 보기</a></td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c_1u7JSBMx3x"
      },
      "source": [
        "# 사용자 정의 미분\n",
        "\n",
        "이 튜토리얼은 사용자 정의 파생물을 정의하고, 파생 수술을 수행하고, 단 5줄의 Swift에서 고유한 그래디언트 체크포인트 API를 구현하는 방법을 보여줍니다."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gHuQo_kCTjFr"
      },
      "source": [
        "## 사용자 정의 파생물 선언하기"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LP0gMw56TlvH"
      },
      "source": [
        "미분 가능한 매개변수와 결과가 있는 모든 Swift 함수에 대해 사용자 정의 파생물을 정의할 수 있습니다. 그렇게 함으로써 C 함수를 가져와 미분할 수도 있습니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j0a8prgZTlEO"
      },
      "outputs": [],
      "source": [
        "import Glibc\n",
        "\n",
        "func sillyExp(_ x: Float) -> Float {\n",
        "    let 𝑒 = Float(M_E)\n",
        "    print(\"Taking 𝑒(\\(𝑒)) to the power of \\(x)!\")\n",
        "    return pow(𝑒, x)\n",
        "}\n",
        "\n",
        "@derivative(of: sillyExp)\n",
        "func sillyDerivative(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {\n",
        "    let y = sillyExp(x)\n",
        "    return (value: y, pullback: { v in v * y })\n",
        "}\n",
        "\n",
        "print(\"exp(3) =\", sillyExp(3))\n",
        "print(\"𝛁exp(3) =\", gradient(of: sillyExp)(3))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eQPX9r3R5OP-"
      },
      "source": [
        "## 파생물 전파 중지하기\n",
        "\n",
        "머신러닝 사용 사례에서 일반적으로 '그래디언트 중지하기'라고 알려진 메서드 `withoutDerivative(at:)`은 파생물이 전파되는 것을 막습니다.\n",
        "\n",
        "또한, `withoutDerivative(at:)`은 때때로 Swift 컴파일러가 미분하지 않을 것을 식별하고 더 효율적인 파생물을 생성하는 데 도움을 줄 수 있습니다. 함수의 도함수가 항상 0이라는 것이 감지되면, Swift 컴파일러는 경고를 생성합니다. `withoutDerivative(at:)`을 명시적으로 사용하면 해당 경고가 발생하지 않습니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ctRt6vBO5Wle"
      },
      "outputs": [],
      "source": [
        "let x: Float = 2.0\n",
        "let y: Float = 3.0\n",
        "gradient(at: x, y) { x, y in\n",
        "    sin(sin(sin(x))) + withoutDerivative(at: cos(cos(cos(y))))\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EeV3wXQ79WS2"
      },
      "source": [
        "## 파생 수술\n",
        "\n",
        "[`withDerivative(_:)`](https://www.tensorflow.org/swift/api_docs/Protocols/Differentiable#/s:10TensorFlow14DifferentiablePAAE12withGradientyxy15CotangentVectorQzzcF) 메서드는 둘러싸는 함수의 역전파 동안 일정 값을 사용해 그래디언트에서 임의의 연산(변경 포함)을 실행합니다.\n",
        "\n",
        "이를 사용하여 역전파를 디버깅하거나 실험적으로 조정해 보세요."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AHV0ryTiD6j8"
      },
      "source": [
        "### 어디서나 동작합니다"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9zKSeUjTmbxq"
      },
      "source": [
        "표준 라이브러리에서 제공하는 모든 미분 API는 `Differentiable` 프로토콜을 준수하는 모든 형식에 대해 일반적으로 정의되며, 이에는 `Float`, `Double`, `Float80`, SIMD 벡터 및 사용자 고유 형식도 포함됩니다!\n",
        "\n",
        "<code>Differentiable</code> 프로토콜에 대한 자세한 내용은 기술 문서 <a>미분 가능 형식</a>을 참조하세요."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eKne7szjD8lr"
      },
      "outputs": [],
      "source": [
        "var x: Float = 30\n",
        "gradient(at: x) { x -> Float in\n",
        "    // Print the partial derivative with respect to the result of `sin(x)`.\n",
        "    let a = sin(x).withDerivative { print(\"∂+/∂sin = \\($0)\") } \n",
        "    // Force the partial derivative with respect to `x` to be `0.5`.\n",
        "    let b = log(x.withDerivative { (dx: inout Float) in\n",
        "        print(\"∂log/∂x = \\(dx), but rewritten to 0.5\");\n",
        "        dx = 0.5\n",
        "    })\n",
        "    return a + b\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vmw0gkqlD9xf"
      },
      "source": [
        "### 신경망 모듈에서 사용하기"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JCf_OplsWzhW"
      },
      "source": [
        "간단한 `Float` 함수에서 사용한 것과 마찬가지로 [TensorFlow Deep Learning 라이브러리용 Swift](https://github.com/tensorflow/swift-apis)를 사용하여 구축된 다음 신경망과 같은 모든 숫자 애플리케이션에서 사용할 수 있습니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fnSeAbs9-hf3"
      },
      "outputs": [],
      "source": [
        "import TensorFlow\n",
        "\n",
        "struct MLP: Layer {\n",
        "    var layer1 = Dense<Float>(inputSize: 2, outputSize: 10, activation: relu)\n",
        "    var layer2 = Dense<Float>(inputSize: 10, outputSize: 1, activation: relu)\n",
        "    \n",
        "    @differentiable\n",
        "    func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {\n",
        "        let h0 = layer1(input).withDerivative { print(\"∂L/∂layer1 =\", $0) }\n",
        "        return layer2(h0)\n",
        "    }\n",
        "}\n",
        "\n",
        "var classifier = MLP()\n",
        "let optimizer = SGD(for: classifier, learningRate: 0.02)\n",
        "\n",
        "let x: Tensor<Float> = [[0, 0], [0, 1], [1, 0], [1, 1]]\n",
        "let y: Tensor<Float> = [0, 1, 1, 0]\n",
        "\n",
        "for _ in 0..<10 {\n",
        "    let 𝛁model = gradient(at: classifier) { classifier -> Tensor<Float> in\n",
        "        let ŷ = classifier(x).withDerivative { print(\"∂L/∂ŷ =\", $0) }\n",
        "        let loss = (ŷ - y).squared().mean()\n",
        "        print(\"Loss: \\(loss)\")\n",
        "        return loss\n",
        "    }\n",
        "    optimizer.update(&classifier, along: 𝛁model)\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TzLfTj28gEUD"
      },
      "source": [
        "## 메모리를 절약하기 위해 역전파 중 활성화 다시 계산하기(체크포인트하기)\n",
        "\n",
        "체크포인트는 메모리를 절약하기 위한 역방향 모드 자동 미분의 전통적인 기술입니다. 파생물 계산을 위해 원래 계산에서 큰 중간값을 저장하는 대신 역전파 중에 필요에 따라 중간값이 다시 계산됩니다.\n",
        "\n",
        "이 기술은 현대 딥 러닝 라이브러리에서도 구현되었습니다. Swift에서 API [`withRecomputationInPullbacks(_:)`](https://www.tensorflow.org/swift/api_docs/Protocols/Differentiable#/s:10TensorFlow14DifferentiablePAAE28withRecomputationInPullbacksyqd__qd__xcAaBRd__lF)를 사용하면 역전파 중에 재계산할 항목을 제어할 수 있으며, 모든 `Differentiable` 형식에서 사용할 수 있습니다.\n",
        "\n",
        "그러나 오늘은 몇 줄의 코드만으로 처음부터 자체 그래디언트 체크포인트 API를 정의하는 방법을 알아보겠습니다."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5cZe-JbjwMfZ"
      },
      "source": [
        "### 그래디언트 체크포인트 API"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "606ob1dn2v77"
      },
      "source": [
        "표준 라이브러리 함수 <a><code>differentiableFunction(from:)</code></a> 측면에서 고유의 그래디언트 체크포인트 API인 <code>makeRecomputedInGradient(_:)</code>를 정의할 수 있습니다. 이는 파생 함수에서 직접 미분 함수를 만드는 속기입니다( 'vector-Jacobian products(VJP) 함수'라고도 함).\n",
        "\n",
        "앞서 살펴본 것처럼, 미분 함수는 원래 함수의 결과와 풀백 클로저의 튜플을 반환합니다. `value:`에서 `original(x)`를 반환하고 `original`에서 `pullback(at:in:)`을 호출하여 다시 원래의 함수를 평가하고 풀백을 얻습니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b1uU3tcVwl_1"
      },
      "outputs": [],
      "source": [
        "/// Given a differentiable function, returns the same differentiable function except when\n",
        "/// derivatives of this function are being computed. In that case, values in the original function needed\n",
        "/// for computing the derivatives will be recomputed, instead of being captured by the differential or pullback.\n",
        "///\n",
        "/// - Parameter body: The body of the differentiable function.\n",
        "/// - Returns: The same differentiable function whose derivatives, when computed, will recompute\n",
        "///   some values from the original function.\n",
        "func makeRecomputedInGradient<T: Differentiable, U: Differentiable>(\n",
        "    _ original: @escaping @differentiable (T) -> U\n",
        ") -> @differentiable (T) -> U {\n",
        "    return differentiableFunction { x in\n",
        "        (value: original(x), pullback: { v in pullback(at: x, in: original)(v) })\n",
        "    }\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UbeKj7NEF7zz"
      },
      "source": [
        "### 동작하는지 확인하기"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oee8SXital45"
      },
      "outputs": [],
      "source": [
        "let input: Float = 10.0\n",
        "print(\"Running original computation...\")\n",
        "\n",
        "// Differentiable multiplication with checkpointing.\n",
        "let square = makeRecomputedInGradient { (x: Float) -> Float in\n",
        "    print(\"  Computing square...\")\n",
        "    return x * x\n",
        "}\n",
        "\n",
        "// Differentiate `f(x) = (cos(x))^2`.\n",
        "let (output, backprop) = valueWithPullback(at: input) { input -> Float in\n",
        "    return square(cos(input))\n",
        "}\n",
        "print(\"Running backpropagation...\")\n",
        "let grad = backprop(1)\n",
        "print(\"Gradient = \\(grad)\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7SxWsSUqF9Bh"
      },
      "source": [
        "### 신경망 모듈로 확장하기\n",
        "\n",
        "이 예에서는 간단한 컨볼루셔널 신경망을 정의합니다.\n",
        "\n",
        "```swift\n",
        "struct Model: Layer {     var conv = Conv2D<Float>(filterShape: (5, 5, 3, 6))     var maxPool = MaxPool2D<Float>(poolSize: (2, 2), strides: (2, 2))     var flatten = Flatten<Float>()     var dense = Dense<Float>(inputSize: 36 * 6, outputSize: 10)      @differentiable     func call(_ input: Tensor<Float>) -> Tensor<Float> {         return input.sequenced(through: conv, maxPool, flatten, dense)     } }\n",
        "```\n",
        "\n",
        "역전파 중 컨볼루션 레이어 (`conv`)의 활성화를 다시 계산해야 합니다. 그러나 `makeRecomputedInGradient(_:)`를 사용하면, 특히 [`sequenced(in:through:_:_:_:_:)`](https://www.tensorflow.org/swift/api_docs/Protocols/Differentiable#/s:10TensorFlow14DifferentiablePAAE9sequenced2in7through____6OutputQyd_3_AA7ContextC_qd__qd_0_qd_1_qd_2_qd_3_t5InputQyd__RszAA5LayerRd__AaMRd_0_AaMRd_1_AaMRd_2_AaMRd_3_AKQyd_0_AGRtd__AKQyd_1_AGRtd_0_AKQyd_2_AGRtd_1_AKQyd_3_AGRtd_2_r3_lF)를 사용하여 레이어를 순차적으로 적용하려는 경우, 결과 코드가 번거롭게 보일 수 있습니다.\n",
        "\n",
        "```swift\n",
        "input.sequenced(in: context, through: conv, maxPool, flatten, dense)\n",
        "```\n",
        "\n",
        "그렇다면 레이어를 래핑하고 역전파 중에 활성화를 다시 계산하는 **특수 레이어 형식**을 정의해 보겠습니다."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZP86M5RjP3OG"
      },
      "source": [
        "먼저 이진 함수를 받는 `makeRecomputedInGradient(_:)` 함수를 정의합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bEm-n5H0QB8s"
      },
      "outputs": [],
      "source": [
        "// Same as the previous `makeRecomputedInGradient(_:)`, except it's for binary functions.\n",
        "func makeRecomputedInGradient<T: Differentiable, U: Differentiable, V: Differentiable>(\n",
        "    _ original: @escaping @differentiable (T, U) -> V\n",
        ") -> @differentiable (T, U) -> V {\n",
        "    return differentiableFunction { x, y in\n",
        "        (value: original(x, y), pullback: { v in pullback(at: x, y, in: original)(v) })\n",
        "    }\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YU6DgqXxP5Nl"
      },
      "source": [
        "그런 다음 제네릭 레이어 `ActivationDiscarding<Wrapped>`를 정의합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ao1r_lIPGeOl"
      },
      "outputs": [],
      "source": [
        "import TensorFlow\n",
        "\n",
        "/// A layer wrapper that makes the underlying layer's activations be discarded during application\n",
        "/// and recomputed during backpropagation.\n",
        "struct ActivationDiscarding<Wrapped: Layer>: Layer {\n",
        "    /// The wrapped layer.\n",
        "    var wrapped: Wrapped\n",
        "\n",
        "    @differentiable\n",
        "    func callAsFunction(_ input: Wrapped.Input) -> Wrapped.Output {\n",
        "        let apply = makeRecomputedInGradient { (layer: Wrapped, input: Input) -> Wrapped.Output in\n",
        "            print(\"    Applying \\(Wrapped.self) layer...\")\n",
        "            return layer(input)\n",
        "        }\n",
        "        return apply(wrapped, input)\n",
        "    }\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HqPXwwuTRjmz"
      },
      "source": [
        "마지막으로 같은 레이어를 반환하는 모든 레이어에 메서드를 추가할 수 있습니다. 단, 적용 중에는 활성화가 삭제되고 역전파 중에 다시 계산됩니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PGgkNnNNR1th"
      },
      "outputs": [],
      "source": [
        "extension Layer {\n",
        "    func discardingActivations() -> ActivationDiscarding<Self> {\n",
        "        return ActivationDiscarding(wrapped: self)\n",
        "    }\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8PP-NZ9XU5_n"
      },
      "source": [
        "모델로 돌아가서 컨볼루션 레이어를 활성화-재계산 레이어로 래핑하면 마무리됩니다.\n",
        "\n",
        "```swift\n",
        "var conv = Conv2D<Float>(filterShape: (5, 5, 3, 6)).discardingActivations()\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bCwNPtCfSbGi"
      },
      "source": [
        "이제 모델에서 사용해 보세요!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gsWGwFjOJ3Md"
      },
      "outputs": [],
      "source": [
        "struct Model: Layer {\n",
        "    var conv = Conv2D<Float>(filterShape: (5, 5, 3, 6)).discardingActivations()\n",
        "    var maxPool = MaxPool2D<Float>(poolSize: (2, 2), strides: (2, 2))\n",
        "    var flatten = Flatten<Float>()\n",
        "    var dense = Dense<Float>(inputSize: 36 * 6, outputSize: 10)\n",
        "\n",
        "    @differentiable\n",
        "    func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {\n",
        "        return input.sequenced(through: conv, maxPool, flatten, dense)\n",
        "    }\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dmFxciU6VYdF"
      },
      "source": [
        "훈련 루프를 실행하면 컨볼루션 레이어의 활성화가 레이어 적용 중에 한 번, 역전파 중에 한 번, 총 두 번 계산되는 것을 볼 수 있습니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-x1nYu0uVSPn"
      },
      "outputs": [],
      "source": [
        "// Use random training data.\n",
        "let x = Tensor<Float>(randomNormal: [10, 16, 16, 3])\n",
        "let y = Tensor<Int32>(rangeFrom: 0, to: 10, stride: 1)\n",
        "\n",
        "var model = Model()\n",
        "let opt = SGD(for: model)\n",
        "\n",
        "for i in 1...5 {\n",
        "    print(\"Starting training step \\(i)\")\n",
        "    print(\"  Running original computation...\")\n",
        "    let (logits, backprop) = model.appliedForBackpropagation(to: x)\n",
        "    let (loss, dL_dŷ) = valueWithGradient(at: logits) { logits in\n",
        "        softmaxCrossEntropy(logits: logits, labels: y)\n",
        "    }\n",
        "    print(\"  Loss: \\(loss)\")\n",
        "    print(\"  Running backpropagation...\")\n",
        "    let (dL_dθ, _) = backprop(dL_dŷ)\n",
        "    \n",
        "    opt.update(&model, along: dL_dθ)\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gzRaZLa_WX0u"
      },
      "source": [
        "보시다시피 다양한 도메인에 대해 제네릭 미분 가능한 프로그래밍 라이브러리를 정의하는 것은 매우 쉽습니다."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "custom_differentiation.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Swift",
      "name": "swift"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
